import gym
from gym import spaces
import numpy as np

class ProcessDecisionEnv(gym.Env):
    def __init__(self, log_data):
        super(ProcessDecisionEnv, self).__init__()
        
        # 定义状态空间（示例：流程节点+资源负载）
        self.observation_space = spaces.Dict({
            "node": spaces.Discrete(10),  # 10个流程节点
            "resource_util": spaces.Box(low=0, high=1, shape=(3,)),  # 3类资源利用率
            "time_elapsed": spaces.Box(low=0, high=100, shape=(1,))  # 已耗时
        })

        
        # 定义动作空间（路径选择）
        self.action_space = spaces.Discrete(3)  # 3种路径选项
        
        # 加载历史日志数据
        self.log_data = log_data
        self.current_idx = 0
        
    def reset(self):
        # 重置环境并返回初始状态
        self.current_idx = 0
        return self._get_state(0)
    
    def step(self, action):
        # 执行动作并返回 (next_state, reward, done, info)
        next_idx = self.current_idx + 1
        state = self._get_state(next_idx)
        
        # 计算奖励（目标：最小化耗时和资源成本）
        reward = - (state["time_elapsed"][0] * 0.1 + np.mean(state["resource_util"]) * 0.3)
        
        # 判断是否终止（流程结束或超时）
        done = next_idx >= len(self.log_data) or state["time_elapsed"][0] > 50
        
        self.current_idx = next_idx
        return state, reward, done, {}
    
    def _get_state(self, idx):
        # 从日志数据提取状态（实际项目需自定义解析逻辑）
        return {
            "node": self.log_data[idx]["node"],
            "resource_util": np.array(self.log_data[idx]["resource_util"]),
            "time_elapsed": np.array([self.log_data[idx]["time_elapsed"]])
        }